#!/usr/bin/env python
# -*- coding: utf-8 -*-


from __future__ import print_function
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from scipy.spatial.transform import Rotation


# Part of the code is referred from: https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py

def quat2mat(quat):
	x, y, z, w = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]

	B = quat.size(0)

	w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
	wx, wy, wz = w*x, w*y, w*z
	xy, xz, yz = x*y, x*z, y*z

	rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
						  2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
						  2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).reshape(B, 3, 3)
	return rotMat

def pose2mat(quat, translation):

	# [R t] [R1 t1] = [R*R1 R*t1+t]
	# [0 1] [ 0  1] = [   0    1  ]

	x, y, z, w = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]

	B = quat.size(0)

	w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
	wx, wy, wz = w*x, w*y, w*z
	xy, xz, yz = x*y, x*z, y*z

	t1, t2, t3 = translation[:, 0], translation[:, 1], translation[:, 2]
	l0 = torch.constant(0.0)
	l1 = torch.constant(1.0)

	mat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz, t1,
						  2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx, t2,
						  2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2, t3,
						  l0, l0, l0, l1], dim=1).reshape(B, 4, 4)
	return mat


def transform_point_cloud(point_cloud, rotation, translation):
	if len(rotation.size()) == 2:
		rot_mat = quat2mat(rotation)
	else:
		rot_mat = rotation
	return torch.matmul(rot_mat, point_cloud) + translation.unsqueeze(2)


def combine_transformations(rotation, translation, rotation_prev, translation_prev):
	return torch.matmul(rotation, rotation_prev), (torch.matmul(rotation, translation_prev.unsqueeze(2)) + translation.unsqueeze(2))[:,:,0]

def npmat2euler(mats, seq='zyx'):
	eulers = []
	for i in range(mats.shape[0]):
		r = Rotation.from_dcm(mats[i])
		eulers.append(r.as_euler(seq, degrees=True))
	return np.asarray(eulers, dtype='float32')